"""
AI Agent Planner 实现
使用 LangChain 框架构建智能客服 Agent
"""

import json
import logging
from typing import Dict, Any, Optional
from langchain.agents import initialize_agent, AgentType
from langchain.prompts import PromptTemplate
from langchain.schema import BaseLanguageModel

from .tools import OCRTool, QueryOrderTool, GenerateRefundLinkTool
from .prompt import SYSTEM_PROMPT, USER_PROMPT_TEMPLATE

# 配置日志
logger = logging.getLogger(__name__)


class CustomerServiceAgent:
    """智能客服 Agent 类"""

    def __init__(self, llm: BaseLanguageModel):
        """
        初始化客服 Agent

        Args:
            llm: 语言模型实例
        """
        self.llm = llm
        self.tools = self._initialize_tools()
        self.agent = self._create_agent()
        logger.info("Customer Service Agent initialized")

    def _initialize_tools(self):
        """初始化工具列表"""
        return [
            OCRTool(),
            QueryOrderTool(),
            GenerateRefundLinkTool()
        ]

    def _create_agent(self):
        """创建 LangChain Agent"""
        try:
            agent = initialize_agent(
                tools=self.tools,
                llm=self.llm,
                agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # 使用零样本推理
                verbose=True,
                max_iterations=3,  # 限制最大迭代次数
                max_execution_time=30,  # 限制执行时间
                handle_parsing_errors=True,  # 处理解析错误
                system_message=SYSTEM_PROMPT
            )
            return agent
        except Exception as e:
            logger.error(f"Failed to create agent: {e}")
            raise

    def process_request(self, user_input: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        处理用户请求

        Args:
            user_input: 用户输入
            context: 额外的上下文信息（如OCR结果）

        Returns:
            处理结果字典
        """
        try:
            # 构建完整的输入
            full_input = user_input
            if context:
                if 'ocr_result' in context:
                    full_input += f"\n[OCR结果]{json.dumps(context['ocr_result'], ensure_ascii=False)}"

            logger.info(f"Processing request: {user_input[:100]}...")

            # 调用 Agent
            response = self.agent.run(full_input)

            # 解析响应
            result = self._parse_response(response)

            logger.info(f"Request processed successfully")
            return result

        except Exception as e:
            logger.error(f"Error processing request: {e}")
            return self._create_error_response(str(e))

    def _parse_response(self, response: str) -> Dict[str, Any]:
        """
        解析 Agent 响应

        Args:
            response: Agent 原始响应

        Returns:
            结构化响应字典
        """
        try:
            # 尝试解析 JSON 响应
            if response.strip().startswith('{'):
                parsed = json.loads(response)
                return {
                    "thought": parsed.get("thought", "处理中..."),
                    "action": parsed.get("action", ""),
                    "observation": parsed.get("observation", ""),
                    "answer": parsed.get("answer", response),
                    "success": True
                }
            else:
                # 非 JSON 响应，直接返回
                return {
                    "thought": "直接回复",
                    "action": "",
                    "observation": "",
                    "answer": response,
                    "success": True
                }

        except json.JSONDecodeError:
            # JSON 解析失败
            return {
                "thought": "解析响应",
                "action": "",
                "observation": "",
                "answer": response,
                "success": True
            }

    def _create_error_response(self, error_msg: str) -> Dict[str, Any]:
        """创建错误响应"""
        return {
            "thought": "处理请求时发生错误",
            "action": "",
            "observation": f"错误: {error_msg}",
            "answer": f"抱歉，处理您的请求时出现问题：{error_msg}。请稍后重试或联系客服。",
            "success": False
        }

    async def aprocess_request(self, user_input: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        异步处理用户请求

        Args:
            user_input: 用户输入
            context: 额外的上下文信息

        Returns:
            处理结果字典
        """
        # 这里可以实现异步版本
        # 目前先调用同步版本
        return self.process_request(user_input, context)


def get_customer_service_agent(llm: BaseLanguageModel) -> CustomerServiceAgent:
    """
    获取智能客服 Agent 实例

    Args:
        llm: 语言模型实例

    Returns:
        CustomerServiceAgent 实例
    """
    return CustomerServiceAgent(llm)
